{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6d1ca9ac",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/bedrock.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e3a8796-edc8-43f2-94ad-fe4fb20d70ed",
   "metadata": {},
   "source": [
    "# Bedrock"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b007403c-6b7a-420c-92f1-4171d05ed9bb",
   "metadata": {},
   "source": [
    "## Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ead155e-b8bd-46f9-ab9b-28fc009361dd",
   "metadata": {},
   "source": [
    "#### Call `complete` with a prompt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "211db735",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81732d83",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60be18ae-c957-4ac2-a58a-0652e18ee6d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import Bedrock\n",
    "\n",
    "resp = Bedrock(\n",
    "    model=\"amazon.titan-text-express-v1\",\n",
    ").complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac2cbebb-a444-4a46-9d85-b265a3483d68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "son of John and Mary Graham.\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14831268-f90f-499d-9d86-925dbc88292b",
   "metadata": {},
   "source": [
    "#### Call `chat` with a list of messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe29574-4af1-48d5-9739-f60652b6ce6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import ChatMessage, Bedrock\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"Tell me a story\"),\n",
    "]\n",
    "\n",
    "resp = Bedrock(\n",
    "    model=\"amazon.titan-text-express-v1\",\n",
    ").chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cbd550a-0264-4a11-9b2c-a08d8723a5ae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant:  Ahoy, matey! Here's a tale for you:\n",
      "\n",
      "Once upon a time, in a faraway land, there was a pirate named Captain Jack Sparrow. He was the most notorious pirate on the seven seas, with a reputation for his wit, charm, and daring exploits.\n",
      "\n",
      "Captain Jack Sparrow was born on the island of Tortuga, where he grew up learning the ways of the sea from his father, Captain Teague Sparrow. As a young boy, Jack was fascinated by the stories of pirates and their adventures, and he dreamed of one day setting sail on his own ship.\n",
      "\n",
      "One day, when Jack was just a teenager, his father was captured by the Spanish Navy. Jack was determined to rescue his father, and he set out on a quest to find the Spanish galleon that had taken him.\n",
      "\n",
      "After a long and treacherous journey, Jack finally found the Spanish galleon. He boarded the ship and fought his way to his father's cell, where he rescued him and escaped the ship with his father.\n",
      "\n",
      "From that day on, Jack Sparrow became a pirate, sailing the high seas and seeking his fortune. He encountered many strange and exotic creatures, including mermaids, sea monsters, and even the legendary Kraken.\n",
      "\n",
      "One of Captain Jack Sparrow's most famous adventures was his encounter with the Black Pearl, a legendary pirate ship that was said to be cursed. Captain Jack Sparrow stole the Black Pearl from a group of Spanish pirates and set out on a quest to find the legendary Fountain of Youth.\n",
      "\n",
      "Along the way, Jack encountered many challenges, including a group of pirates led by Captain Barbossa, who was determined to reclaim the Black Pearl. Captain Jack Sparrow and his crew fought bravely against Barbossa's pirates, and in the end, they were able to reclaim the Black Pearl and escape with their lives.\n",
      "\n",
      "Captain Jack Sparrow was a skilled navigator and a master of disguise. He was able to fool many of his enemies, including the Spanish Navy, by using his wit and charm to outsmart them.\n",
      "\n",
      "In the end, Captain Jack Sparrow was a legendary pirate, and his adventures will be remembered for generations to come. He was a true hero, and his legacy will continue to inspire people to seek their fortune and adventure on the high seas.\n",
      "\n",
      "Would you like to hear more tales of adventure and bravery?\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ed5e894-4597-4911-a623-591560f72b82",
   "metadata": {},
   "source": [
    "## Streaming"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cb7986f-aaed-42e2-abdd-f274f6d4fc59",
   "metadata": {},
   "source": [
    "Using `stream_complete` endpoint "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d43f17a2-0aeb-464b-a7a7-732ba5e8ef24",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import Bedrock\n",
    "\n",
    "llm = Bedrock(model=\"amazon.titan-text-express-v1\")\n",
    "resp = llm.stream_complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0214e911-cf0d-489c-bc48-9bb1d8bf65d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "son of John and Mary Graham."
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40350dd8-3f50-4a2f-8545-5723942039bb",
   "metadata": {},
   "source": [
    "Using `stream_chat` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc636e65-a67b-4dcd-ac60-b25abc9d8dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import Bedrock\n",
    "\n",
    "llm = Bedrock(\n",
    "    model=\"amazon.titan-text-express-v1\",\n",
    ")\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"Tell me a story\"),\n",
    "]\n",
    "resp = llm.stream_chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4475a6bc-1051-4287-abce-ba83324aeb9e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Ahoy, matey! Here's a tale for you:\n",
      "\n",
      "Once upon a time, there was a pirate named Captain Jack Sparrow. He was a notorious pirate on the high seas, known for his wit, his charm, and his love of adventure.\n",
      "\n",
      "Captain Jack Sparrow was born on the Caribbean island of Tortuga. His father, Captain Teague Sparrow, was a notorious pirate himself, and Jack grew up learning the ways of the sea from his father.\n",
      "\n",
      "At a young age, Jack showed a natural talent for sailing and navigation. He could navigate his ship through the toughest storms and find his way through the most treacherous waters.\n",
      "\n",
      "One day, Jack's father was captured by the Spanish Navy. The Spanish were determined to capture all the pirates on the Caribbean, and they were particularly interested in Captain Teague Sparrow.\n",
      "\n",
      "Jack was determined to rescue his father. He gathered a crew of his most trusted friends, and they set out on a daring mission to rescue Captain Teague Sparrow from the Spanish prison ship.\n",
      "\n",
      "The mission was perilous, and the crew faced many challenges. They had to navigate through treacherous waters, avoid the Spanish Navy, and deal with a host of other pirates and sea creatures.\n",
      "\n",
      "But Jack was a skilled captain, and he led his crew through the challenges with courage and determination. They finally reached the Spanish prison ship, and Jack and his crew attacked the guards.\n",
      "\n",
      "They managed to rescue Captain Teague Sparrow, but they were soon surrounded by the Spanish Navy. Jack and his crew fought bravely, but they were outnumbered and outgunned.\n",
      "\n",
      "Just when it seemed like all was lost, Jack had an idea. He used his wit and charm to distract the Spanish Navy, while his crew sneaked aboard their ship.\n",
      "\n",
      "They managed to steal the Spanish ship's treasure, and they set sail for Tortuga, their home island.\n",
      "\n",
      "Captain Jack Sparrow and his crew were hailed as heroes on Tortuga. They were celebrated for their bravery and their cunning, and they were known as the most feared pirates on the high seas.\n",
      "\n",
      "But Jack was not content with just being a hero. He was determined to find the legendary treasure of the Spanish Main. He knew that it was hidden somewhere on the ocean, and he was determined to find it.\n",
      "\n",
      "Jack and his crew set out on a new adventure, and they encountered many challenges along the way. They faced storms,"
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "009d3f1c-ef35-4126-ae82-0b97adb746e3",
   "metadata": {},
   "source": [
    "## Configure Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e973e3d1-a3c9-43b9-bee1-af3e57946ac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import Bedrock\n",
    "\n",
    "llm = Bedrock(\n",
    "    model=\"amazon.titan-text-express-v1\", temperature=0.5, max_tokens=2048\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2c9bcf6-c950-4dfc-abdc-598d5bdedf40",
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = llm.complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de715fea-4878-4fbb-b415-71250215f3a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "an English singer, songwriter, and guitarist who is widely regarded as one of the most influential and important musicians of the 20th century. He is best known for his work with the Beatles, which he joined in 1962, and for his solo career, which began in 1965.\n",
      "Paul McCartney is known for his unique songwriting style, which combines elements of pop, rock, classical, and folk music. He is also known for his versatile guitar playing, which has inspired many musicians throughout the world.\n",
      "\n",
      "One of Paul McCartney's most famous songs is \"Yesterday,\" which he wrote and performed with the Beatles in 1965. The song has become one of the most covered songs in history and has been praised for its beautiful melody and clever lyrics.\n",
      "\n",
      "Paul McCartney is also known for his activism and humanitarian work. He has been a champion for environmental causes, peace, and social justice throughout his career, and has used his platform to raise awareness about these issues. He has been awarded numerous honors, including the Nobel Peace Prize in 2019, for his contributions to music and society.\n",
      "\n",
      "Paul McCartney was born on September 18, 1942, in Liverpool, England. He grew up in a working-class family and showed an early interest in music. He began playing the guitar at the age of 12 and formed his first band, the Quarrymen, in 1957. The Quarrymen later became the Beatles, and Paul McCartney became the band's lead singer and guitarist.\n",
      "\n",
      "The Beatles' first hit song, \"Love Me Do,\" was released in 1962, and the band quickly became one of the most popular bands in the world. The Beatles' music was characterized by its catchy melodies, innovative production techniques, and clever lyrics, and the band's popularity continued to grow throughout the 1960s.\n",
      "\n",
      "Paul McCartney's solo career began in 1965 with the release of his debut album, \"McCartney.\" The album was a huge success, and Paul McCartney became known as a solo artist in his own right. He continued to release successful albums throughout the 1970s, including \"Band on the Run,\" \"Let It Be,\" and \"Abbey Road.\"\n",
      "\n",
      "In addition to his music career, Paul McCartney has been involved in many other ventures. He has written books, directed films, and even designed a line of clothing. He is also a member of the band Wings, which he formed in 1971 with his wife Linda McCartney.\n",
      "\n",
      "Paul McCartney is widely regarded as one of the greatest musicians of all time. His unique songwriting style, versatile guitar playing, and activism have made him a beloved figure in the music industry and a respected figure in society. He continues to perform and record new music, and his legacy as a musician and humanitarian will continue to inspire generations to come.\n",
      "\n",
      "Paul McCartney is an English singer, songwriter, and guitarist who is widely regarded as one of the most influential and important musicians of the 20th century. He is best known for his work with the Beatles, which he joined in 1962, and for his solo career, which began in 1965.\n",
      "\n",
      "One of Paul McCartney's most famous songs is \"Yesterday,\" which he wrote and performed with the Beatles in 1965. The song has become one of the most covered songs in history and has been praised for its beautiful melody and clever lyrics.\n",
      "\n",
      "Paul McCartney is also known for his activism and humanitarian work. He has been a champion for environmental causes, peace, and social justice throughout his career, and has used his platform to raise awareness about these issues. He has been awarded numerous honors, including the Nobel Peace Prize in 2019, for his contributions to music and society.\n",
      "\n",
      "Paul McCartney was born on September 18, 1942, in Liverpool, England. He grew up in a working-class family and showed an early interest in music. He began playing the guitar at the age of 12 and formed his first band, the Quarrymen, in 1957. The Quarrymen later became the Beatles, and Paul McCartney became the band's lead singer and guitarist.\n",
      "\n",
      "The Beatles' first hit song, \"Love Me Do,\" was released in 1962, and the band quickly became one of the most popular bands in the world. The Beatles' music was characterized by its catchy melodies, innovative production techniques, and clever lyrics, and the band's popularity continued to grow throughout the 1960s.\n",
      "\n",
      "Paul McCartney's solo career began in 1965 with the release of his debut album, \"McCartney.\" The album was a huge success, and Paul McCartney became known as a solo artist in his own right. He continued to release successful albums throughout the 1970s, including \"Band on the Run,\" \"Let It Be,\" and \"Abbey Road.\"\n",
      "\n",
      "In addition to his music career, Paul McCartney has been involved in many other ventures. He has written books, directed films, and even designed a line of clothing. He is also a member of the band Wings, which he formed in 1971 with his wife Linda McCartney.\n",
      "\n",
      "Paul McCartney is widely regarded as one of the greatest musicians of all time. His unique songwriting style, versatile guitar playing, and activism have made him a beloved figure in the music industry and a respected figure in society. He continues to perform and record new music, and his legacy as a musician and humanitarian will continue to inspire generations to come.\n",
      "\n",
      "Paul McCartney is an English singer, songwriter, and guitarist who is widely regarded as one of the most influential and important musicians of the 20th century. He is best known for his work with the Beatles, which he joined in 1962, and for his solo career, which began in 1965.\n",
      "\n",
      "One of Paul McCartney's most famous songs is \"Yesterday,\" which he wrote and performed with the Beatles in 1965. The song has become one of the most covered songs in history and has been praised for its beautiful melody and clever lyrics.\n",
      "\n",
      "Paul McCartney is also known for his activism and humanitarian work. He has been a champion for environmental causes, peace, and social justice throughout his career, and has used his platform to raise awareness about these issues. He has been awarded numerous honors, including the Nobel Peace Prize in 2019, for his contributions to music and society.\n",
      "\n",
      "Paul McCartney was born on September 18, 1942, in Liverpool, England. He grew up in a working-class family and showed an early interest in music. He began playing the guitar at the age of 12 and formed his first band, the Quarrymen, in 1957. The Quarrymen later became the Beatles, and Paul McCartney became the band's lead singer and guitarist.\n",
      "\n",
      "The Beatles' first hit song, \"Love Me Do,\" was released in 1962, and the band quickly became one of the most popular bands in the world. The Beatles' music was characterized by its catchy melodies, innovative production techniques, and clever lyrics, and the band's popularity continued to grow throughout the 1960s.\n",
      "\n",
      "Paul McCartney's solo career began in 1965 with the release of his debut album, \"McCartney.\" The album was a huge success, and Paul McCartney became known as a solo artist in his own right. He continued to release successful albums throughout the 1970s, including \"Band on the Run,\" \"Let It Be,\" and \"Abbey Road.\"\n",
      "\n",
      "In addition to his music career, Paul McCartney has been involved in many other ventures. He has written books, directed films, and even designed a line of clothing. He is also a member of the band Wings, which he formed in 1971 with his wife Linda McCartney.\n",
      "\n",
      "Paul McCartney is widely regarded as one of the greatest musicians of all time. His unique songwriting style, versatile guitar playing, and activism have made him a beloved figure in the music industry and a respected figure in society. He continues to perform and record new music, and his legacy as a musician and humanitarian will continue to inspire generations to come.\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bdd4602-e37c-4230-af82-35af5292f9a0",
   "metadata": {},
   "source": [
    "# Connect to Bedrock with Access Keys "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9c80814-6d59-4782-a4bb-cbfcdba6a072",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import Bedrock\n",
    "\n",
    "llm = Bedrock(\n",
    "    model=\"amazon.titan-text-express-v1\",\n",
    "    aws_access_key_id=\"AWS Access Key ID to use\",\n",
    "    aws_secret_access_key=\"AWS Secret Access Key to use\",\n",
    "    aws_session_token=\"AWS Session Token to use\",\n",
    "    region_name=\"AWS Region to use, eg. us-east-1\",\n",
    "    profile_name=\"AWS Profile to use (if not using aws_access_key_id and aws_secret_access_key)\",\n",
    ")\n",
    "\n",
    "resp = llm.complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07982f4-035c-41ca-9dca-49ae6ab3c05a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "son of John and Mary Graham.\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
